// CodeInjectionDlg.cpp : implementation file
//

#include "stdafx.h"
#include "CodeInjection.h"
#include "CodeInjectionDlg.h"
#include "afxdialogex.h"
#include <string>

#include <keystone/keystone.h>

#include <sstream>
#include <iostream>
#include <iomanip>

#include <vector>

#ifdef _DEBUG
#define new DEBUG_NEW
#endif

// CAboutDlg dialog used for App About

class CAboutDlg : public CDialogEx
{
public:
	CAboutDlg();

	// Dialog Data
#ifdef AFX_DESIGN_TIME
	enum { IDD = IDD_ABOUTBOX };
#endif

protected:
	virtual void DoDataExchange(CDataExchange* pDX);    // DDX/DDV support

// Implementation
protected:
	DECLARE_MESSAGE_MAP()
};

CAboutDlg::CAboutDlg() : CDialogEx(IDD_ABOUTBOX)
{
}

void CAboutDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialogEx::DoDataExchange(pDX);
}

BEGIN_MESSAGE_MAP(CAboutDlg, CDialogEx)
END_MESSAGE_MAP()

// CCodeInjectionDlg dialog

CCodeInjectionDlg::CCodeInjectionDlg(CWnd* pParent /*=NULL*/)
	: CDialogEx(IDD_CODEINJECTION_DIALOG, pParent)
{
	m_hIcon = AfxGetApp()->LoadIcon(IDR_MAINFRAME);
}

void CCodeInjectionDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialogEx::DoDataExchange(pDX);
	DDX_Control(pDX, IDC_EDIT_ASM, _edit_asm);
	DDX_Control(pDX, IDC_EDIT_OPCODE, _edit_opcode);
	DDX_Control(pDX, IDC_CHECK_ISX64, _checkbox_isx64);
}

BEGIN_MESSAGE_MAP(CCodeInjectionDlg, CDialogEx)
	ON_WM_SYSCOMMAND()
	ON_WM_PAINT()
	ON_WM_QUERYDRAGICON()
	ON_BN_CLICKED(IDC_BUTTON_TRANSFER, &CCodeInjectionDlg::OnBnClickedButtonTransfer)
END_MESSAGE_MAP()

// CCodeInjectionDlg message handlers

BOOL CCodeInjectionDlg::OnInitDialog()
{
	CDialogEx::OnInitDialog();

	// Add "About..." menu item to system menu.

	// IDM_ABOUTBOX must be in the system command range.
	ASSERT((IDM_ABOUTBOX & 0xFFF0) == IDM_ABOUTBOX);
	ASSERT(IDM_ABOUTBOX < 0xF000);

	CMenu* pSysMenu = GetSystemMenu(FALSE);
	if (pSysMenu != NULL)
	{
		BOOL bNameValid;
		CString strAboutMenu;
		bNameValid = strAboutMenu.LoadString(IDS_ABOUTBOX);
		ASSERT(bNameValid);
		if (!strAboutMenu.IsEmpty())
		{
			pSysMenu->AppendMenu(MF_SEPARATOR);
			pSysMenu->AppendMenu(MF_STRING, IDM_ABOUTBOX, strAboutMenu);
		}
	}

	// Set the icon for this dialog.  The framework does this automatically
	//  when the application's main window is not a dialog
	SetIcon(m_hIcon, TRUE);			// Set big icon
	SetIcon(m_hIcon, FALSE);		// Set small icon

	// TODO: Add extra initialization here

	return TRUE;  // return TRUE  unless you set the focus to a control
}

void CCodeInjectionDlg::OnSysCommand(UINT nID, LPARAM lParam)
{
	if ((nID & 0xFFF0) == IDM_ABOUTBOX)
	{
		CAboutDlg dlgAbout;
		dlgAbout.DoModal();
	}
	else
	{
		CDialogEx::OnSysCommand(nID, lParam);
	}
}

// If you add a minimize button to your dialog, you will need the code below
//  to draw the icon.  For MFC applications using the document/view model,
//  this is automatically done for you by the framework.

void CCodeInjectionDlg::OnPaint()
{
	if (IsIconic())
	{
		CPaintDC dc(this); // device context for painting

		SendMessage(WM_ICONERASEBKGND, reinterpret_cast<WPARAM>(dc.GetSafeHdc()), 0);

		// Center icon in client rectangle
		int cxIcon = GetSystemMetrics(SM_CXICON);
		int cyIcon = GetSystemMetrics(SM_CYICON);
		CRect rect;
		GetClientRect(&rect);
		int x = (rect.Width() - cxIcon + 1) / 2;
		int y = (rect.Height() - cyIcon + 1) / 2;

		// Draw the icon
		dc.DrawIcon(x, y, m_hIcon);
	}
	else
	{
		CDialogEx::OnPaint();
	}
}

int CCodeInjectionDlg::getOpcode(
	ks_arch arch,
	int mode,
	const char *assembly,
	int syntax,
	OUT unsigned char *opcode,
	OUT unsigned int *psize
)
{
	ks_engine *ks;
	ks_err err;
	size_t count;
	unsigned char *encode;
	size_t size;

	err = ks_open(arch, mode, &ks);
	if (err != KS_ERR_OK)
	{
		//printf("ERROR: failed on ks_open(), quit\n");
		return -1;
	}

	if (syntax)
		ks_option(ks, KS_OPT_SYNTAX, syntax);

	if (ks_asm(ks, assembly, 0, &encode, &size, &count))
	{
		//printf("ERROR: failed on ks_asm() with count = %lu, error code = %u\n", count, ks_errno(ks));
	}
	else
	{
		size_t i;

		if (psize)
		{
			*psize = size;
		}

		//printf("%s = ", assembly);
		for (i = 0; i < size; i++)
		{
			opcode[i] = encode[i];
			//printf("%02x ", encode[i]);
		}

		//opcode += ss.str();

		//printf("\n");
		//printf("Assembled: %lu bytes, %lu statements\n\n", size, count);
	}

	// NOTE: free encode after usage to avoid leaking memory
	ks_free(encode);

	// close Keystone instance when done
	ks_close(ks);

	return 0;
}
// The system calls this function to obtain the cursor to display while the user drags
//  the minimized window.
HCURSOR CCodeInjectionDlg::OnQueryDragIcon()
{
	return static_cast<HCURSOR>(m_hIcon);
}

void CCodeInjectionDlg::OnBnClickedButtonTransfer()
{
	// TODO: Add your control notification handler code here

	CString csAsm;
	_edit_asm.GetWindowTextW(csAsm);
	//AfxMessageBox(csAsm);

	std::wstring wline(L"");
	//auto x = csAsm.GetLength();

	std::vector<std::wstring> v;
	v.clear();

	for (auto i(0l); i < csAsm.GetLength(); ++i)
	{
		if (csAsm[i] == L'\r' && csAsm[i + 1] == '\n'

			)
		{
			v.push_back(wline);
			wline = L"";
			continue;
		}
		wline += csAsm[i];
	}

	v.push_back(wline);

	//CString cs;
	//cs.Format(L"%d", x);
	std::auto_ptr<BYTE> ap{ new BYTE[1024]{} };
	unsigned int size(0);
	std::stringstream ss("");
	int mode;

	try
	{
		if (_checkbox_isx64.GetCheck() == BST_CHECKED)
		{
			//X64
			mode = KS_MODE_64;
		}
		else
		{
			mode = KS_MODE_32;
		}

		for (auto i : v)
		{
			std::string sasm(i.begin(), i.end());
			//MessageBoxA(0, sasm.data(), 0, 0);

			if (getOpcode(KS_ARCH_X86, mode, sasm.data(), 0, ap.get(), &size) == -1)
			{
				continue;
			}
			for (auto j(0ul); j < size; ++j)
			{
				ss << std::setw(2) << std::setfill('0') << std::hex << (DWORD)ap.get()[j] << " ";
			}

			ss << "\r\n";
			auto ssss(ss.str());
			std::wstring str(ssss.begin(), ssss.end());
			_edit_opcode.SetWindowTextW(str.data());
			//MessageBoxA(0, ss.str().data(), 0, 0);
		}
	}
	catch (const std::exception&)
	{
	}
}